{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DeepMSim: A Deep Learning based MS/MS spectra Simulator\n",
    "### *Muhammad Haseeb and Sumesh Kumar*\n",
    "##### *CAP5610: Machine Learning - Fall 2019*\n",
    "##### *School of Computing and Information Sciences*\n",
    "##### *Florida International University (FIU)*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Licence Information"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This Transformer model construction for our project has been derived from the official <a href=\"https://tensorflow.org\" class=\"external\">Tensorflow</a> <a href=\"https://github.com/tensorflow/docs/blob/master/site/en/tutorials/text/transformer.ipynb\" class=\"external\">Transformer tutorial</a>."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The MIT License\n",
    "\n",
    "#### Copyright 2019 Muhammad Haseeb and Sumesh Kumar\n",
    "\n",
    "Permission is hereby granted, free of charge, to any person obtaining a copy of this software \n",
    "and associated documentation files (the \"Software\"), to deal in the Software without restriction, \n",
    "including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, \n",
    "and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, \n",
    "subject to the following conditions:\n",
    "\n",
    "The above copyright notice and this permission notice shall be included in all copies or substantial \n",
    "portions of the Software.\n",
    "\n",
    "THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT \n",
    "LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. \n",
    "IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, \n",
    "WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE \n",
    "SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.\n",
    "\n",
    "### Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "\n",
    "you may not use this file except in compliance with the License.\n",
    "You may obtain a copy of the License [here](https://www.apache.org/licenses/LICENSE-2.0).\n",
    "    \n",
    "Unless required by applicable law or agreed to in writing, software\n",
    "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "See the License for the specific language governing permissions and\n",
    "limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Transformer model for MS/MS spectra simulation\n",
    "We will implement a <a href=\"https://arxiv.org/abs/1706.03762\" class=\"external\">Transformer model</a> to translate the peptide sequences into spectral peak sequences using TensorFlow.\n",
    "\n",
    "Transformer model employs the concept of *self-attention* — the ability to attend to different positions of the input sequence to compute a representation of that sequence for modelling. It is composed of several layer multiconnected encoders and decoders that construct the output sequences. \n",
    "\n",
    "However, the drawback is that if the input sequences have a *temporal/spatial* (positional) relationships such as in peptide sequence data, a positional encoding such as <a href=\"https://arxiv.org/abs/1810.04805BERT\" class=\"external\">BERT encoding</a> must be added or the transformers will only see the input sequences as bag {set} of unordered and unrelated words. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import the required libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Required Imports\n",
    "from __future__ import absolute_import, division, print_function, unicode_literals\n",
    "\n",
    "try:\n",
    "  %tensorflow_version 2.x\n",
    "except Exception:\n",
    "  pass\n",
    "import tensorflow_datasets as tfds\n",
    "import tensorflow as tf\n",
    "import pandas as pd\n",
    "import os\n",
    "import re\n",
    "import sys\n",
    "import math\n",
    "import time\n",
    "import urllib\n",
    "import numpy as np\n",
    "import operator\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.model_selection import train_test_split"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Acquisition\n",
    "Let's get the annotated MS/MS spectral data libraries from <a href=\"https://chemdata.nist.gov/dokuwiki/doku.php?id=peptidew:cdownload\" class=\"external\">NIST.</a>\n",
    "\n",
    "**CAUTION:** Make sure the downloaded spectral library datasets have been sourced from the same Mass-spectrometer instrument using the same fragmentation method.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# HYPERPARAMETER: Choose between HUMAN and MOUSE.\n",
    "# species = HUMAN\n",
    "# species = MOUSE\n",
    "\n",
    "species = 'MOUSE'\n",
    "\n",
    "assert (species == 'HUMAN' or species == 'MOUSE')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "if species == 'HUMAN':\n",
    "    human = 'cptac2_human_hcd_itraq_selected_part1.msp.tar.gz'\n",
    "    msp_file = './speclib/cptac2_human_hcd_itraq_selected_part1.msp'\n",
    "\n",
    "    data_dir = os.path.expanduser(\"./speclib\")\n",
    "    if not os.path.isdir(data_dir):\n",
    "        os.mkdir(data_dir)\n",
    "\n",
    "    if not os.path.exists('./speclib/cptac2_human_hcd_itraq_selected_part1.msp'):\n",
    "        if not os.path.exists(human):\n",
    "            !wget 'ftp://chemdata.nist.gov/download/peptide_library/libraries/cptaclib/2015/cptac2_human_hcd_itraq_selected_part1.msp.tar.gz'\n",
    "    \n",
    "        # Extract the tar ball into MSP file and move into the speclib folder\n",
    "        !tar xvf 'cptac2_human_hcd_itraq_selected_part1.msp.tar.gz'\n",
    "        os.rename('./cptac2_human_hcd_itraq_selected_part1.msp', msp_file)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Getting the Mouse Orbitrap spectral library (.MSP format) from NIST and placing it in the speclib folder\n",
    "if species == 'MOUSE':\n",
    "    mouse = 'cptac2_mouse_hcd_selected.msp.tar.gz'\n",
    "    msp_file = './speclib/cptac2_mouse_hcd_selected.msp'\n",
    "\n",
    "    data_dir = os.path.expanduser(\"./speclib\")\n",
    "    if not os.path.isdir(data_dir):\n",
    "        os.mkdir(data_dir)\n",
    "\n",
    "    if not os.path.exists('./speclib/cptac2_mouse_hcd_selected.msp'):\n",
    "        if not os.path.exists(mouse):\n",
    "            !wget 'ftp://chemdata.nist.gov/download/peptide_library/libraries/cptaclib/2015/cptac2_mouse_hcd_selected.msp.tar.gz'\n",
    "\n",
    "        # Extract the tar ball into MSP file and move into the speclib folder\n",
    "        !tar xvf 'cptac2_mouse_hcd_selected.msp.tar.gz'\n",
    "        os.rename('./cptac2_mouse_hcd_selected.msp', msp_file)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Extraction and Preprocessing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Mine the spectral libraries and extract corpuses for both languages. i.e. peptide sequences and their MS/MS spectra.\n",
    "Construct a dataframe with sentences and their translations, with words separated by a '*space* and ending with ' .' for both languages. Then, Convert the dataframe to a tensor dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Process known Immonium Ions\n",
    "The commonly observed immonium ions along with their isotopes are added to a dictionary un order to represent the ion sentences (from the corpus) in terms of the dictionary (bag of words)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Known immonium ions\n",
    "imm = np.array(['AA', 'RA', 'RB', 'RC', 'RD', 'RE', 'RF', 'RG', 'RJ', 'NA', 'NB', 'DA', 'DB', 'CA', 'EA', 'QA', 'QB',\n",
    "                'QC', 'QD', 'GA', 'HA', 'HB', 'HC', 'HD', 'HE', 'HF', 'IA', 'IB', 'IC', 'LA', 'LB', 'LC', 'KA', 'KB',\n",
    "                'KC', 'KD', 'KE', 'KF', 'MA', 'MB', 'FA', 'FB', 'PA', 'SA', 'TA', 'WA', 'WB', 'WC', 'WD', 'WE', 'WF',\n",
    "                'WG', 'YA', 'YB', 'YC', 'VA', 'VB', 'VC', 'VD'])\n",
    "\n",
    "# TODO: Add isotopes?\n",
    "'''\n",
    "# Make isotopes of an array of ions\n",
    "iso = []\n",
    "\n",
    "# For each immonium ion, append the +i and +2i isotopes\n",
    "for im in imm:\n",
    "    iso.append(im + '+i')\n",
    "    iso.append(im + '+2i')\n",
    "\n",
    "# Append the isotopes to the main dictionary\n",
    "imm = np.append(imm, np.array(iso))\n",
    "'''\n",
    "# Make a set representation\n",
    "simm= set(imm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Corpus Extraction\n",
    "Extract the peptide sequence and ion sentence sequences from the dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### MSP file parser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create empty lists to store temporary data\n",
    "seq = []\n",
    "X = []\n",
    "Z = []\n",
    "E = []\n",
    "\n",
    "# Parse the MSP file and get the immonium ion dataset\n",
    "def parseMSPfile(fname=None, dimm=None):\n",
    "    header = False\n",
    "    stnc = []\n",
    "    with open(fname, 'r') as slib:\n",
    "        for line in slib:\n",
    "            # Check for empty line for spectrum end and header start\n",
    "            if (line[0] == '\\r' or line[0] == '#' or line[0] == '\\n'):\n",
    "                header = False\n",
    "                X.append(stnc[:])\n",
    "                stnc.clear()\n",
    "                continue\n",
    "\n",
    "            # Check if spectrum header\n",
    "            if (header == False):\n",
    "                if (line[:7] != 'Comment'):\n",
    "                    key, val = line.split(':')\n",
    "\n",
    "                    if key == 'Name':\n",
    "                        # Take out the sequence\n",
    "                        nm, chg = val.split('/')\n",
    "                        seq.append(nm[1:])\n",
    "                        # Take out the charge, nmods and eV (energy)\n",
    "                        ch = chg.split('_')\n",
    "\n",
    "                        # Append the charge\n",
    "                        Z.append(int(ch[0]))\n",
    "                        \n",
    "                        # Append the energy\n",
    "                        #print(ch)\n",
    "                        E.append(float(ch[-1][:-3]))\n",
    "                        \n",
    "\n",
    "                    if key == 'Num peaks':\n",
    "                        header = True\n",
    "\n",
    "            # If not the header, then it is the peak list\n",
    "            else:\n",
    "                _, _, label = line.split('\\t')\n",
    "\n",
    "                # Clip off the first and last \" \"\n",
    "                label = label[1:-1]\n",
    "\n",
    "                # Extract label pattern\n",
    "                lbls = re.findall(r\"(?<!\\/)(?<![A-Z])(I[A-Z][A-Z](?![A-Z]){2}[\\/])\", label)\n",
    "                \n",
    "                # FIXME: To also find isotopes\n",
    "                #re.findall(r\"(?<!\\/)(?<![A-Z])(I[A-Z][A-Z](?![A-Z]){2}[+]*\\d*i*)\", label)\n",
    "                \n",
    "                lbls = list(map(lambda x: x[1:], lbls))\n",
    "                \n",
    "                lbls = list(map(lambda x: x[:-1], lbls))\n",
    "\n",
    "                # If anything to append then append\n",
    "                if lbls:\n",
    "                    for lab in lbls:\n",
    "                        if lab not in dimm:\n",
    "                            dimm.add(lab[0])\n",
    "                            #print ('Label Added:' + lab)\n",
    "                    #print(lbls)\n",
    "                    stnc.extend(lbls)\n",
    "\n",
    "    # For the last one\n",
    "    X.append(stnc[:])\n",
    "    stnc.clear()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clear the lists\n",
    "seq = []\n",
    "X = []\n",
    "Z = []\n",
    "E = []\n",
    "\n",
    "# Parse the MSP file and create a dataframe\n",
    "_ = parseMSPfile(fname = msp_file, dimm = simm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Construct the two languages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# filter data based on the max charge\n",
    "MAXZ = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                              seq  z     E              ions\n",
      "0          AAAAAAAAAAAAAAAGAGAGAK  2  47.5          [KD, KF]\n",
      "1          AAAAAAAAAAAAAAAGAGAGAK  2  66.0              [KF]\n",
      "2          AAAAAAAAAAAAAAAGAGAGAK  3  30.0              [KF]\n",
      "3          AAAAAAAAAAAAAAAGAGAGAK  3  33.0              [KF]\n",
      "4          AAAAAAAAAAAAAAAGAGAGAK  3  37.0              [KF]\n",
      "6                     AAAAAAAAAAR  2  26.0  [RG, RI, RA, RJ]\n",
      "7   AAAAAAAAAGEAINPTYYDCHFAPEVGGR  3  60.0              [HA]\n",
      "8                      AAAAAAAAAK  2  37.0              [KF]\n",
      "9   AAAAAAAAAPAAAATAPTTAATTAATAAQ  2  63.4                []\n",
      "10  AAAAAAAAAPAAAATAPTTAATTAATAAQ  3  39.9          [QA, QD]\n",
      "12                   AAAAAAAAAVSR  2  28.9  [RG, RI, RA, RJ]\n",
      "13                   AAAAAAAAAVSR  2  32.0          [RG, RJ]\n",
      "14                   AAAAAAAAAVSR  2  40.0      [RG, RI, RJ]\n",
      "15                   AAAAAAAAAVSR  3  18.2  [RG, RI, RA, RJ]\n",
      "16      AAAAAAAAVPSAGPAGPAPTSAAGR  3  34.6              [RJ]\n",
      "17                    AAAAAAALQAK  2  31.4              [QA]\n",
      "18                    AAAAAAALQAK  2  35.0              [QA]\n",
      "19                    AAAAAAALQAK  2  36.0                []\n",
      "20                    AAAAAAALQAK  2  41.0              [QA]\n",
      "21                    AAAAAAALQAK  2  43.0              [QA]\n"
     ]
    }
   ],
   "source": [
    "# Construct a dataframe\n",
    "dataset = pd.DataFrame(list(zip(seq, Z, E, X)), columns=['seq', 'z', 'E', 'ions'])\n",
    "\n",
    "# Filter the dataset by the max charge \n",
    "dataset = dataset[dataset['z'] <= MAXZ]\n",
    "\n",
    "# Print a few entries to visualize\n",
    "print(dataset[:20])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sentence Structurization\n",
    "\n",
    "The ions sentences must be aligned to their peptide sequences in order to establish the notion of positional encoding for the two languages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preProcess(x,y):\n",
    "    if y:\n",
    "        n = len(y)\n",
    "        indices = []\n",
    "        result = []\n",
    "        for i in range(n):\n",
    "            indices.append(x.index(y[i][0]))\n",
    "        tuples = list(zip(y,indices))\n",
    "        tuples.sort(key=operator.itemgetter(1))\n",
    "        result,second = zip(*tuples)\n",
    "        result = \" \".join(result) + ' -'\n",
    "        return result\n",
    "    else:\n",
    "        return ' -'\n",
    "        \n",
    "\n",
    "\n",
    "lan1 = list(map(lambda x: \" \".join(x), dataset['seq']))\n",
    "lan2 = list(map(preProcess, dataset['seq'], dataset['ions']))\n",
    "\n",
    "# Clear the duplicate data\n",
    "seq.clear()\n",
    "X.clear()\n",
    "Z.clear()\n",
    "E.clear()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Insert the languages back to the dataframe\n",
    "dataset['seq'] = lan1\n",
    "dataset['ions'] = lan2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's encode the charge and energy inside \n",
    "# the input peptide sequence sentence\n",
    "\n",
    "# Convert z and E to strings\n",
    "dataset['z'] = dataset['z'].astype(str)\n",
    "dataset['E'] = dataset['E'].astype(str)\n",
    "\n",
    "# TODO: Create a new column named enc with everything included in sentences\n",
    "dataset['enc'] = dataset['z'] + ' ' + dataset['seq'] + ' ' + dataset['E'] + ' -' #+  dataset['E'] + ' -'  + ' ' +"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Language Encoding\n",
    "\n",
    "Tokenize the sentences, construct tensors, embed the charge and eV information and do positional encoding (BERT)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Perform the 80-20 train-test split\n",
    "train, test = train_test_split(dataset, test_size=0.2)\n",
    "\n",
    "# Create a tf dataset object\n",
    "train_examples = tf.data.Dataset.from_tensor_slices((train['enc'],train['ions']))\n",
    "test_examples = tf.data.Dataset.from_tensor_slices((test['enc'],test['ions']))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sanity Checks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(b'3 G T Y T T D D S P S D I A E I R L D K 37.9 -', shape=(), dtype=string)\n",
      "tf.Tensor(b'YA EA RJ KF -', shape=(), dtype=string)\n",
      "tf.Tensor(b'2 A V I E M S S K 36.0 -', shape=(), dtype=string)\n",
      "tf.Tensor(b'KF -', shape=(), dtype=string)\n",
      "tf.Tensor(b'2 A V A W E A G K 28.2 -', shape=(), dtype=string)\n",
      "tf.Tensor(b'WD WE WA WF EA KD KF -', shape=(), dtype=string)\n",
      "tf.Tensor(b'3 F S D L D Y M G M S S G F Y K 32.3 -', shape=(), dtype=string)\n",
      "tf.Tensor(b'FA YC YA MA KF -', shape=(), dtype=string)\n",
      "tf.Tensor(b'3 E A A T A V E Q E K 21.7 -', shape=(), dtype=string)\n",
      "tf.Tensor(b'EA QA QD KF -', shape=(), dtype=string)\n",
      "<class 'tensorflow.python.data.ops.dataset_ops.TensorSliceDataset'>\n"
     ]
    }
   ],
   "source": [
    "cnt = 0\n",
    "for pep, ions in train_examples:\n",
    "    print (pep)\n",
    "    print (ions)\n",
    "    cnt+=1\n",
    "    if (cnt>=5):\n",
    "        break\n",
    "        \n",
    "print(type(train_examples))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tokenization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer_pep = tfds.features.text.SubwordTextEncoder.build_from_corpus(\n",
    "    (pep.numpy() for pep, ions in train_examples), target_vocab_size=2**10)\n",
    "\n",
    "tokenizer_ion = tfds.features.text.SubwordTextEncoder.build_from_corpus(\n",
    "    (ions.numpy() for pep, ions in train_examples), target_vocab_size=2**10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sanity Check"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tokenized string is [15, 11, 1, 11, 12, 10, 6, 1, 72, 179, 184, 4]\n",
      "The original string: 3 P E P T I D E 64.3 -\n",
      "KFRG RJ -\n",
      "15 ----> 3 \n",
      "11 ----> P \n",
      "1 ----> E \n",
      "11 ----> P \n",
      "12 ----> T \n",
      "10 ----> I \n",
      "6 ----> D \n",
      "1 ----> E \n",
      "72 ----> 64\n",
      "179 ----> .\n",
      "184 ----> 3\n",
      "4 ---->  -\n"
     ]
    }
   ],
   "source": [
    "sample_string = '3 P E P T I D E 64.3 -'\n",
    "\n",
    "tokenized_string = tokenizer_pep.encode(sample_string)\n",
    "print ('Tokenized string is {}'.format(tokenized_string))\n",
    "\n",
    "original_string = tokenizer_pep.decode(tokenized_string)\n",
    "print ('The original string: {}'.format(original_string))\n",
    "\n",
    "\n",
    "tok = np.array([5,   4,   2,   1])\n",
    "print(tokenizer_ion.decode(tok))\n",
    "\n",
    "assert original_string == sample_string\n",
    "\n",
    "for ts in tokenized_string:\n",
    "  print ('{} ----> {}'.format(ts, tokenizer_pep.decode([ts])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hyperparameters\n",
    "BUFFER_SIZE = 20000\n",
    "BATCH_SIZE = 64"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Add a start and end token to the input and target."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def encode(lang1, lang2):\n",
    "  lang1 = [tokenizer_pep.vocab_size] + tokenizer_pep.encode(\n",
    "      lang1.numpy()) + [tokenizer_pep.vocab_size+1]\n",
    "\n",
    "  lang2 = [tokenizer_ion.vocab_size] + tokenizer_ion.encode(\n",
    "      lang2.numpy()) + [tokenizer_ion.vocab_size+1]\n",
    "  \n",
    "  return lang1, lang2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Operations inside `.map()` run in graph mode and receive a graph tensor that do not have a numpy attribute. The `tokenizer` expects a string or Unicode symbol to encode it into integers. Hence, you need to run the encoding inside a `tf.py_function`, which receives an eager tensor having a numpy attribute that contains the string value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tf_encode(pep, ions):\n",
    "  return tf.py_function(encode, [pep, ions], [tf.int64, tf.int64])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = train_examples.map(tf_encode)\n",
    "\n",
    "# TODO: Do we need this filter?\n",
    "#train_dataset = train_dataset.filter(filter_max_length)\n",
    "\n",
    "# cache the dataset to memory to get a speedup while reading from it.\n",
    "train_dataset = train_dataset.cache()\n",
    "train_dataset = train_dataset.shuffle(BUFFER_SIZE).padded_batch(\n",
    "    BATCH_SIZE, padded_shapes=([-1], [-1]))\n",
    "train_dataset = train_dataset.prefetch(tf.data.experimental.AUTOTUNE)\n",
    "\n",
    "# The test dataset\n",
    "test_dataset = test_examples.map(tf_encode)\n",
    "test_dataset = test_dataset.padded_batch(\n",
    "    BATCH_SIZE, padded_shapes=([-1], [-1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Positional encoding\n",
    "\n",
    "Since this model doesn't contain any recurrence or convolution, positional encoding is added to give the model some information about the relative position of the words in the sentence. \n",
    "\n",
    "The positional encoding vector is added to the embedding vector. Embeddings represent a token in a d-dimensional space where tokens with similar meaning will be closer to each other. But the embeddings do not encode the relative position of words in a sentence. So after adding the positional encoding, words will be closer to each other based on the *similarity of their meaning and their position in the sentence*, in the d-dimensional space. The formula for calculating the positional encoding is as follows:\n",
    "\n",
    "$$\\Large{PE_{(pos, 2i)} = sin(pos / 10000^{2i / d_{model}})} $$\n",
    "$$\\Large{PE_{(pos, 2i+1)} = cos(pos / 10000^{2i / d_{model}})} $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_angles(pos, i, d_model):\n",
    "  angle_rates = 1 / np.power(10000, (2 * (i//2)) / np.float32(d_model))\n",
    "  return pos * angle_rates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def positional_encoding(position, d_model):\n",
    "  angle_rads = get_angles(np.arange(position)[:, np.newaxis],\n",
    "                          np.arange(d_model)[np.newaxis, :],\n",
    "                          d_model)\n",
    "  \n",
    "  # apply sin to even indices in the array; 2i\n",
    "  angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])\n",
    "  \n",
    "  # apply cos to odd indices in the array; 2i+1\n",
    "  angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])\n",
    "    \n",
    "  pos_encoding = angle_rads[np.newaxis, ...]\n",
    "    \n",
    "  return tf.cast(pos_encoding, dtype=tf.float32)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Masking"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Mask all the pad tokens in the batch of sequence. It ensures that the model does not treat padding as the input. The mask indicates where pad value `0` is present: it outputs a `1` at those locations, and a `0` otherwise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_padding_mask(seq):\n",
    "  seq = tf.cast(tf.math.equal(seq, 0), tf.float32)\n",
    "  \n",
    "  # add extra dimensions to add the padding\n",
    "  # to the attention logits.\n",
    "  return seq[:, tf.newaxis, tf.newaxis, :]  # (batch_size, 1, 1, seq_len)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_look_ahead_mask(size):\n",
    "  mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0)\n",
    "  return mask  # (seq_len, seq_len)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Scaled dot product attention"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The attention function used by the transformer takes three inputs: Q (query), K (key), V (value). The equation used to calculate the attention weights is:\n",
    "\n",
    "$$\\Large{Attention(Q, K, V) = softmax_k(\\frac{QK^T}{\\sqrt{d_k}}) V} $$\n",
    "\n",
    "The dot-product attention is scaled by a factor of square root of the depth. This is done because for large values of depth, the dot product grows large in magnitude pushing the softmax function where it has small gradients resulting in a very hard softmax. \n",
    "\n",
    "For example, consider that `Q` and `K` have a mean of 0 and variance of 1. Their matrix multiplication will have a mean of 0 and variance of `dk`. Hence, *square root of `dk`* is used for scaling (and not any other number) because the matmul of `Q` and `K` should have a mean of 0 and variance of 1, and you get a gentler softmax.\n",
    "\n",
    "The mask is multiplied with -1e9 (close to negative infinity). This is done because the mask is summed with the scaled matrix multiplication of Q and K and is applied immediately before a softmax. The goal is to zero out these cells, and large negative inputs to softmax are near zero in the output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "def scaled_dot_product_attention(q, k, v, mask):\n",
    "  \"\"\"Calculate the attention weights.\n",
    "  q, k, v must have matching leading dimensions.\n",
    "  k, v must have matching penultimate dimension, i.e.: seq_len_k = seq_len_v.\n",
    "  The mask has different shapes depending on its type(padding or look ahead) \n",
    "  but it must be broadcastable for addition.\n",
    "  \n",
    "  Args:\n",
    "    q: query shape == (..., seq_len_q, depth)\n",
    "    k: key shape == (..., seq_len_k, depth)\n",
    "    v: value shape == (..., seq_len_v, depth_v)\n",
    "    mask: Float tensor with shape broadcastable \n",
    "          to (..., seq_len_q, seq_len_k). Defaults to None.\n",
    "    \n",
    "  Returns:\n",
    "    output, attention_weights\n",
    "  \"\"\"\n",
    "\n",
    "  matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)\n",
    "  \n",
    "  # scale matmul_qk\n",
    "  dk = tf.cast(tf.shape(k)[-1], tf.float32)\n",
    "  scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)\n",
    "\n",
    "  # add the mask to the scaled tensor.\n",
    "  if mask is not None:\n",
    "    scaled_attention_logits += (mask * -1e9)  \n",
    "\n",
    "  # softmax is normalized on the last axis (seq_len_k) so that the scores\n",
    "  # add up to 1.\n",
    "  attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)  # (..., seq_len_q, seq_len_k)\n",
    "\n",
    "  output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)\n",
    "\n",
    "  return output, attention_weights"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multi-head attention"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Multi-head attention consists of four parts:\n",
    "*    Linear layers and split into heads.\n",
    "*    Scaled dot-product attention.\n",
    "*    Concatenation of heads.\n",
    "*    Final linear layer."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each multi-head attention block gets three inputs; Q (query), K (key), V (value). These are put through linear (Dense) layers and split up into multiple heads. \n",
    "\n",
    "The `scaled_dot_product_attention` defined above is applied to each head (broadcasted for efficiency). An appropriate mask must be used in the attention step.  The attention output for each head is then concatenated (using `tf.transpose`, and `tf.reshape`) and put through a final `Dense` layer.\n",
    "\n",
    "Instead of one single attention head, Q, K, and V are split into multiple heads because it allows the model to jointly attend to information at different positions from different representational spaces. After the split each head has a reduced dimensionality, so the total computation cost is the same as a single head attention with full dimensionality."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a `MultiHeadAttention` layer to try out. At each location in the sequence, `y`, the `MultiHeadAttention` runs all 8 attention heads across all other locations in the sequence, returning a new vector of the same length at each location."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention(tf.keras.layers.Layer):\n",
    "  def __init__(self, d_model, num_heads):\n",
    "    super(MultiHeadAttention, self).__init__()\n",
    "    self.num_heads = num_heads\n",
    "    self.d_model = d_model\n",
    "    \n",
    "    assert d_model % self.num_heads == 0\n",
    "    \n",
    "    self.depth = d_model // self.num_heads\n",
    "    \n",
    "    self.wq = tf.keras.layers.Dense(d_model)\n",
    "    self.wk = tf.keras.layers.Dense(d_model)\n",
    "    self.wv = tf.keras.layers.Dense(d_model)\n",
    "    \n",
    "    self.dense = tf.keras.layers.Dense(d_model)\n",
    "        \n",
    "  def split_heads(self, x, batch_size):\n",
    "    \"\"\"Split the last dimension into (num_heads, depth).\n",
    "    Transpose the result such that the shape is (batch_size, num_heads, seq_len, depth)\n",
    "    \"\"\"\n",
    "    x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))\n",
    "    return tf.transpose(x, perm=[0, 2, 1, 3])\n",
    "    \n",
    "  def call(self, v, k, q, mask):\n",
    "    batch_size = tf.shape(q)[0]\n",
    "    \n",
    "    q = self.wq(q)  # (batch_size, seq_len, d_model)\n",
    "    k = self.wk(k)  # (batch_size, seq_len, d_model)\n",
    "    v = self.wv(v)  # (batch_size, seq_len, d_model)\n",
    "    \n",
    "    q = self.split_heads(q, batch_size)  # (batch_size, num_heads, seq_len_q, depth)\n",
    "    k = self.split_heads(k, batch_size)  # (batch_size, num_heads, seq_len_k, depth)\n",
    "    v = self.split_heads(v, batch_size)  # (batch_size, num_heads, seq_len_v, depth)\n",
    "    \n",
    "    # scaled_attention.shape == (batch_size, num_heads, seq_len_q, depth)\n",
    "    # attention_weights.shape == (batch_size, num_heads, seq_len_q, seq_len_k)\n",
    "    scaled_attention, attention_weights = scaled_dot_product_attention(\n",
    "        q, k, v, mask)\n",
    "    \n",
    "    scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3])  # (batch_size, seq_len_q, num_heads, depth)\n",
    "\n",
    "    concat_attention = tf.reshape(scaled_attention, \n",
    "                                  (batch_size, -1, self.d_model))  # (batch_size, seq_len_q, d_model)\n",
    "\n",
    "    output = self.dense(concat_attention)  # (batch_size, seq_len_q, d_model)\n",
    "        \n",
    "    return output, attention_weights"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a `MultiHeadAttention` layer to try out. At each location in the sequence, `y`, the `MultiHeadAttention` runs all 8 attention heads across all other locations in the sequence, returning a new vector of the same length at each location."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Point wise feed forward network\n",
    "\n",
    "Point wise feed forward network consists of two fully-connected layers with a ReLU activation in between."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "def point_wise_feed_forward_network(d_model, dff):\n",
    "  return tf.keras.Sequential([\n",
    "      tf.keras.layers.Dense(dff, activation='relu'),  # (batch_size, seq_len, dff)\n",
    "      tf.keras.layers.Dense(d_model)  # (batch_size, seq_len, d_model)\n",
    "  ])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Encoder and decoder\n",
    "\n",
    "The transformer model follows the same general pattern as a standard sequence to sequence with attention model\n",
    "\n",
    "* The input sentence is passed through `N` encoder layers that generates an output for each word/token in the sequence.\n",
    "* The decoder attends on the encoder's output and its own input (self-attention) to predict the next word. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Encoder layer\n",
    "\n",
    "Each encoder layer consists of sublayers:\n",
    "\n",
    "1.   Multi-head attention (with padding mask) \n",
    "2.    Point wise feed forward networks. \n",
    "\n",
    "Each of these sublayers has a residual connection around it followed by a layer normalization. Residual connections help in avoiding the vanishing gradient problem in deep networks.\n",
    "\n",
    "The output of each sublayer is `LayerNorm(x + Sublayer(x))`. The normalization is done on the `d_model` (last) axis. There are N encoder layers in the transformer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderLayer(tf.keras.layers.Layer):\n",
    "  def __init__(self, d_model, num_heads, dff, rate=0.1):\n",
    "    super(EncoderLayer, self).__init__()\n",
    "\n",
    "    self.mha = MultiHeadAttention(d_model, num_heads)\n",
    "    self.ffn = point_wise_feed_forward_network(d_model, dff)\n",
    "\n",
    "    self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
    "    self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
    "    \n",
    "    self.dropout1 = tf.keras.layers.Dropout(rate)\n",
    "    self.dropout2 = tf.keras.layers.Dropout(rate)\n",
    "    \n",
    "  def call(self, x, training, mask):\n",
    "\n",
    "    attn_output, _ = self.mha(x, x, x, mask)  # (batch_size, input_seq_len, d_model)\n",
    "    attn_output = self.dropout1(attn_output, training=training)\n",
    "    out1 = self.layernorm1(x + attn_output)  # (batch_size, input_seq_len, d_model)\n",
    "    \n",
    "    ffn_output = self.ffn(out1)  # (batch_size, input_seq_len, d_model)\n",
    "    ffn_output = self.dropout2(ffn_output, training=training)\n",
    "    out2 = self.layernorm2(out1 + ffn_output)  # (batch_size, input_seq_len, d_model)\n",
    "    \n",
    "    return out2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Decoder layer\n",
    "\n",
    "Each decoder layer consists of sublayers:\n",
    "\n",
    "1.   Masked multi-head attention (with look ahead mask and padding mask)\n",
    "2.   Multi-head attention (with padding mask). V (value) and K (key) receive the *encoder output* as inputs. Q (query) receives the *output from the masked multi-head attention sublayer.*\n",
    "3.   Point wise feed forward networks\n",
    "\n",
    "Each of these sublayers has a residual connection around it followed by a layer normalization. The output of each sublayer is `LayerNorm(x + Sublayer(x))`. The normalization is done on the `d_model` (last) axis.\n",
    "\n",
    "There are N decoder layers in the transformer.\n",
    "\n",
    "As Q receives the output from decoder's first attention block, and K receives the encoder output, the attention weights represent the importance given to the decoder's input based on the encoder's output. In other words, the decoder predicts the next word by looking at the encoder output and self-attending to its own output. See the demonstration above in the scaled dot product attention section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderLayer(tf.keras.layers.Layer):\n",
    "  def __init__(self, d_model, num_heads, dff, rate=0.1):\n",
    "    super(DecoderLayer, self).__init__()\n",
    "\n",
    "    self.mha1 = MultiHeadAttention(d_model, num_heads)\n",
    "    self.mha2 = MultiHeadAttention(d_model, num_heads)\n",
    "\n",
    "    self.ffn = point_wise_feed_forward_network(d_model, dff)\n",
    " \n",
    "    self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
    "    self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
    "    self.layernorm3 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
    "    \n",
    "    self.dropout1 = tf.keras.layers.Dropout(rate)\n",
    "    self.dropout2 = tf.keras.layers.Dropout(rate)\n",
    "    self.dropout3 = tf.keras.layers.Dropout(rate)\n",
    "    \n",
    "    \n",
    "  def call(self, x, enc_output, training, \n",
    "           look_ahead_mask, padding_mask):\n",
    "    # enc_output.shape == (batch_size, input_seq_len, d_model)\n",
    "\n",
    "    attn1, attn_weights_block1 = self.mha1(x, x, x, look_ahead_mask)  # (batch_size, target_seq_len, d_model)\n",
    "    attn1 = self.dropout1(attn1, training=training)\n",
    "    out1 = self.layernorm1(attn1 + x)\n",
    "    \n",
    "    attn2, attn_weights_block2 = self.mha2(\n",
    "        enc_output, enc_output, out1, padding_mask)  # (batch_size, target_seq_len, d_model)\n",
    "    attn2 = self.dropout2(attn2, training=training)\n",
    "    out2 = self.layernorm2(attn2 + out1)  # (batch_size, target_seq_len, d_model)\n",
    "    \n",
    "    ffn_output = self.ffn(out2)  # (batch_size, target_seq_len, d_model)\n",
    "    ffn_output = self.dropout3(ffn_output, training=training)\n",
    "    out3 = self.layernorm3(ffn_output + out2)  # (batch_size, target_seq_len, d_model)\n",
    "    \n",
    "    return out3, attn_weights_block1, attn_weights_block2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Encoder\n",
    "\n",
    "The `Encoder` consists of:\n",
    "1.   Input Embedding\n",
    "2.   Positional Encoding\n",
    "3.   N encoder layers\n",
    "\n",
    "The input is put through an embedding which is summed with the positional encoding. The output of this summation is the input to the encoder layers. The output of the encoder is the input to the decoder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(tf.keras.layers.Layer):\n",
    "  def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size,\n",
    "               maximum_position_encoding, rate=0.1):\n",
    "    super(Encoder, self).__init__()\n",
    "\n",
    "    self.d_model = d_model\n",
    "    self.num_layers = num_layers\n",
    "    \n",
    "    self.embedding = tf.keras.layers.Embedding(input_vocab_size, d_model)\n",
    "    self.pos_encoding = positional_encoding(maximum_position_encoding, \n",
    "                                            self.d_model)\n",
    "    \n",
    "    \n",
    "    self.enc_layers = [EncoderLayer(d_model, num_heads, dff, rate) \n",
    "                       for _ in range(num_layers)]\n",
    "  \n",
    "    self.dropout = tf.keras.layers.Dropout(rate)\n",
    "    \n",
    "  def call(self, x, training, mask):\n",
    "\n",
    "    seq_len = tf.shape(x)[1]\n",
    "    \n",
    "    # adding embedding and position encoding.\n",
    "    x = self.embedding(x)  # (batch_size, input_seq_len, d_model)\n",
    "    x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))\n",
    "    x += self.pos_encoding[:, :seq_len, :]\n",
    "\n",
    "    x = self.dropout(x, training=training)\n",
    "    \n",
    "    for i in range(self.num_layers):\n",
    "      x = self.enc_layers[i](x, training, mask)\n",
    "    \n",
    "    return x  # (batch_size, input_seq_len, d_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Decoder\n",
    "\n",
    "The `Decoder` consists of:\n",
    "1.   Output Embedding\n",
    "2.   Positional Encoding\n",
    "3.   N decoder layers\n",
    "\n",
    "The target is put through an embedding which is summed with the positional encoding. The output of this summation is the input to the decoder layers. The output of the decoder is the input to the final linear layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(tf.keras.layers.Layer):\n",
    "  def __init__(self, num_layers, d_model, num_heads, dff, target_vocab_size,\n",
    "               maximum_position_encoding, rate=0.1):\n",
    "    super(Decoder, self).__init__()\n",
    "\n",
    "    self.d_model = d_model\n",
    "    self.num_layers = num_layers\n",
    "    \n",
    "    self.embedding = tf.keras.layers.Embedding(target_vocab_size, d_model)\n",
    "    self.pos_encoding = positional_encoding(maximum_position_encoding, d_model)\n",
    "    \n",
    "    self.dec_layers = [DecoderLayer(d_model, num_heads, dff, rate) \n",
    "                       for _ in range(num_layers)]\n",
    "    self.dropout = tf.keras.layers.Dropout(rate)\n",
    "    \n",
    "  def call(self, x, enc_output, training, \n",
    "           look_ahead_mask, padding_mask):\n",
    "\n",
    "    seq_len = tf.shape(x)[1]\n",
    "    attention_weights = {}\n",
    "    \n",
    "    x = self.embedding(x)  # (batch_size, target_seq_len, d_model)\n",
    "    x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))\n",
    "    x += self.pos_encoding[:, :seq_len, :]\n",
    "    \n",
    "    x = self.dropout(x, training=training)\n",
    "\n",
    "    for i in range(self.num_layers):\n",
    "      x, block1, block2 = self.dec_layers[i](x, enc_output, training,\n",
    "                                             look_ahead_mask, padding_mask)\n",
    "      \n",
    "      attention_weights['decoder_layer{}_block1'.format(i+1)] = block1\n",
    "      attention_weights['decoder_layer{}_block2'.format(i+1)] = block2\n",
    "    \n",
    "    # x.shape == (batch_size, target_seq_len, d_model)\n",
    "    return x, attention_weights"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Transformer Construction\n",
    "\n",
    "Transformer consists of the encoder, decoder and a final linear layer. The output of the decoder is the input to the linear layer and its output is returned."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Transformer(tf.keras.Model):\n",
    "  def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size, \n",
    "               target_vocab_size, pe_input, pe_target, rate=0.1):\n",
    "    super(Transformer, self).__init__()\n",
    "\n",
    "    self.encoder = Encoder(num_layers, d_model, num_heads, dff, \n",
    "                           input_vocab_size, pe_input, rate)\n",
    "\n",
    "    self.decoder = Decoder(num_layers, d_model, num_heads, dff, \n",
    "                           target_vocab_size, pe_target, rate)\n",
    "\n",
    "    self.final_layer = tf.keras.layers.Dense(target_vocab_size)\n",
    "    \n",
    "  def call(self, inp, tar, training, enc_padding_mask, \n",
    "           look_ahead_mask, dec_padding_mask):\n",
    "\n",
    "    enc_output = self.encoder(inp, training, enc_padding_mask)  # (batch_size, inp_seq_len, d_model)\n",
    "    \n",
    "    # dec_output.shape == (batch_size, tar_seq_len, d_model)\n",
    "    dec_output, attention_weights = self.decoder(\n",
    "        tar, enc_output, training, look_ahead_mask, dec_padding_mask)\n",
    "    \n",
    "    final_output = self.final_layer(dec_output)  # (batch_size, tar_seq_len, target_vocab_size)\n",
    "    \n",
    "    return final_output, attention_weights"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hyperparameter Setting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_layers = 6\n",
    "d_model = 128\n",
    "dff = 512\n",
    "num_heads = 8\n",
    "\n",
    "input_vocab_size = tokenizer_pep.vocab_size + 2\n",
    "target_vocab_size = tokenizer_ion.vocab_size + 2\n",
    "dropout_rate = 0.1\n",
    "MAX_LENGTH = 80\n",
    "\n",
    "# Optimizer Settings\n",
    "beta1 = 0.85\n",
    "beta2 = 0.95\n",
    "alpha = 1e-7"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optimizer\n",
    "\n",
    "Use the Adam optimizer with a custom learning rate scheduler according to the formula in the [paper](https://arxiv.org/abs/1706.03762).\n",
    "\n",
    "$$\\Large{lrate = d_{model}^{-0.5} * min(step{\\_}num^{-0.5}, step{\\_}num * warmup{\\_}steps^{-1.5})}$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):\n",
    "  def __init__(self, d_model, warmup_steps=4000):\n",
    "    super(CustomSchedule, self).__init__()\n",
    "    \n",
    "    self.d_model = d_model\n",
    "    self.d_model = tf.cast(self.d_model, tf.float32)\n",
    "\n",
    "    self.warmup_steps = warmup_steps\n",
    "    \n",
    "  def __call__(self, step):\n",
    "    arg1 = tf.math.rsqrt(step)\n",
    "    arg2 = step * (self.warmup_steps ** -1.5)\n",
    "    \n",
    "    return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_rate = CustomSchedule(d_model)\n",
    "\n",
    "optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=beta1, beta_2=beta2, \n",
    "                                     epsilon=alpha)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Loss Function\n",
    "\n",
    "Since the target sequences are padded, it is important to apply a padding mask when calculating the loss."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_object = tf.keras.losses.SparseCategoricalCrossentropy(\n",
    "    from_logits=True, reduction='none')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss_function(real, pred):\n",
    "  mask = tf.math.logical_not(tf.math.equal(real, 0))\n",
    "  loss_ = loss_object(real, pred)\n",
    "\n",
    "  mask = tf.cast(mask, dtype=loss_.dtype)\n",
    "  loss_ *= mask\n",
    "  \n",
    "  return tf.reduce_mean(loss_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loss = tf.keras.metrics.Mean(name='train_loss')\n",
    "train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(\n",
    "    name='train_accuracy')\n",
    "\n",
    "test_loss = tf.keras.metrics.Mean(name='test_loss')\n",
    "test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(\n",
    "    name='test_accuracy')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "transformer = Transformer(num_layers, d_model, num_heads, dff,\n",
    "                          input_vocab_size, target_vocab_size, \n",
    "                          pe_input=input_vocab_size, \n",
    "                          pe_target=target_vocab_size,\n",
    "                          rate=dropout_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_masks(inp, tar):\n",
    "  # Encoder padding mask\n",
    "  enc_padding_mask = create_padding_mask(inp)\n",
    "  \n",
    "  # Used in the 2nd attention block in the decoder.\n",
    "  # This padding mask is used to mask the encoder outputs.\n",
    "  dec_padding_mask = create_padding_mask(inp)\n",
    "  \n",
    "  # Used in the 1st attention block in the decoder.\n",
    "  # It is used to pad and mask future tokens in the input received by \n",
    "  # the decoder.\n",
    "  look_ahead_mask = create_look_ahead_mask(tf.shape(tar)[1])\n",
    "  dec_target_padding_mask = create_padding_mask(tar)\n",
    "  combined_mask = tf.maximum(dec_target_padding_mask, look_ahead_mask)\n",
    "  \n",
    "  return enc_padding_mask, combined_mask, dec_padding_mask"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create the checkpoint path and the checkpoint manager. This will be used to save checkpoints every `5` epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint_path = \"./checkpoints/train\"\n",
    "\n",
    "ckpt = tf.train.Checkpoint(transformer=transformer,\n",
    "                           optimizer=optimizer)\n",
    "\n",
    "ckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=5)\n",
    "\n",
    "# if a checkpoint exists, restore the latest checkpoint.\n",
    "if ckpt_manager.latest_checkpoint:\n",
    "  ckpt.restore(ckpt_manager.latest_checkpoint)\n",
    "  print ('Latest checkpoint restored!!')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The target is divided into tar_inp and tar_real. tar_inp is passed as an input to the decoder. `tar_real` is that same input shifted by 1: At each location in `tar_input`, `tar_real` contains the next token that should be predicted.\n",
    "\n",
    "For example, `peptide sequence` = \"SOS  P E P T I D E 24.4 EOS\"\n",
    "\n",
    "`tar_inp` =  \"SOS 2 P E P T I D E 24.4\"\n",
    "\n",
    "`tar_real` = \"2 P E P T I D E 24.4 EOS\"\n",
    "\n",
    "The transformer is an auto-regressive model: it makes predictions one part at a time, and uses its output so far to decide what to do next. \n",
    "As the transformer predicts each word, *self-attention* allows it to look at the previous words in the input sequence to better predict the next word.\n",
    "\n",
    "To prevent the model from peaking at the expected output the model uses a look-ahead mask."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "EPOCHS = 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train Step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The @tf.function trace-compiles train_step into a TF graph for faster\n",
    "# execution. The function specializes to the precise shape of the argument\n",
    "# tensors. To avoid re-tracing due to the variable sequence lengths or variable\n",
    "# batch sizes (the last batch is smaller), use input_signature to specify\n",
    "# more generic shapes.\n",
    "\n",
    "train_step_signature = [\n",
    "    tf.TensorSpec(shape=(None, None), dtype=tf.int64),\n",
    "    tf.TensorSpec(shape=(None, None), dtype=tf.int64),\n",
    "]\n",
    "\n",
    "@tf.function(input_signature=train_step_signature)\n",
    "def train_step(inp, tar):\n",
    "  tar_inp = tar[:, :-1]\n",
    "  tar_real = tar[:, 1:]\n",
    "  \n",
    "  enc_padding_mask, combined_mask, dec_padding_mask = create_masks(inp, tar_inp)\n",
    "  \n",
    "  with tf.GradientTape() as tape:\n",
    "    predictions, _ = transformer(inp, tar_inp, \n",
    "                                 True,\n",
    "                                 enc_padding_mask, \n",
    "                                 combined_mask, \n",
    "                                 dec_padding_mask)\n",
    "    loss = loss_function(tar_real, predictions)\n",
    "\n",
    "  gradients = tape.gradient(loss, transformer.trainable_variables)    \n",
    "  optimizer.apply_gradients(zip(gradients, transformer.trainable_variables))\n",
    "  \n",
    "  train_loss(loss)\n",
    "  train_accuracy(tar_real, predictions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test Step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_step_signature = [\n",
    "    tf.TensorSpec(shape=(None, None), dtype=tf.int64),\n",
    "    tf.TensorSpec(shape=(None, None), dtype=tf.int64),\n",
    "]\n",
    "\n",
    "@tf.function(input_signature=test_step_signature)\n",
    "def test_step(inp, tar):\n",
    "  tar_inp = tar[:, :-1]\n",
    "  tar_real = tar[:, 1:]\n",
    "  \n",
    "  enc_padding_mask, combined_mask, dec_padding_mask = create_masks(inp, tar_inp)\n",
    "  \n",
    "  with tf.GradientTape() as tape:\n",
    "    predictions, _ = transformer(inp, tar_inp, \n",
    "                                 True,\n",
    "                                 enc_padding_mask, \n",
    "                                 combined_mask, \n",
    "                                 dec_padding_mask)\n",
    "    loss = loss_function(tar_real, predictions)\n",
    "  \n",
    "  test_loss(loss)\n",
    "  test_accuracy(tar_real, predictions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Peptide Sequence is used as the input language and Ion Sequence is the target language."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1 Batch 0 Loss 0.3329 Accuracy 0.2823\n",
      "Epoch 1 Batch 50 Loss 0.3274 Accuracy 0.3545\n",
      "Epoch 1 Batch 100 Loss 0.3264 Accuracy 0.3539\n",
      "Epoch 1 Batch 150 Loss 0.3289 Accuracy 0.3538\n",
      "Epoch 1 Batch 200 Loss 0.3259 Accuracy 0.3543\n",
      "Epoch 1 Batch 250 Loss 0.3271 Accuracy 0.3556\n",
      "Epoch 1 Batch 300 Loss 0.3267 Accuracy 0.3562\n",
      "Epoch 1 Batch 350 Loss 0.3263 Accuracy 0.3574\n",
      "Epoch 1 Batch 400 Loss 0.3253 Accuracy 0.3574\n",
      "Epoch 1 Batch 450 Loss 0.3260 Accuracy 0.3574\n",
      "Epoch 1 Batch 500 Loss 0.3266 Accuracy 0.3580\n",
      "Epoch 1 Batch 550 Loss 0.3260 Accuracy 0.3578\n",
      "Epoch 1 Batch 600 Loss 0.3254 Accuracy 0.3574\n",
      "Epoch 1 Batch 650 Loss 0.3231 Accuracy 0.3557\n",
      "Epoch 1 Batch 700 Loss 0.3220 Accuracy 0.3551\n",
      "Epoch 1 Batch 750 Loss 0.3222 Accuracy 0.3555\n",
      "Epoch 1 Batch 800 Loss 0.3220 Accuracy 0.3557\n",
      "Epoch 1 Batch 850 Loss 0.3223 Accuracy 0.3559\n",
      "Epoch 1 Batch 900 Loss 0.3215 Accuracy 0.3556\n",
      "Epoch 1 Batch 950 Loss 0.3213 Accuracy 0.3556\n",
      "Epoch 1 Batch 1000 Loss 0.3213 Accuracy 0.3561\n",
      "Epoch 1 Batch 1050 Loss 0.3214 Accuracy 0.3566\n",
      "Epoch 1 Batch 1100 Loss 0.3208 Accuracy 0.3563\n",
      "Epoch 1 Batch 1150 Loss 0.3204 Accuracy 0.3563\n",
      "Epoch 1 Batch 1200 Loss 0.3200 Accuracy 0.3564\n",
      "Epoch 1 Batch 1250 Loss 0.3196 Accuracy 0.3565\n",
      "Epoch 1 Batch 1300 Loss 0.3195 Accuracy 0.3568\n",
      "Epoch 1 Batch 1350 Loss 0.3192 Accuracy 0.3571\n",
      "Epoch 1 Batch 1400 Loss 0.3189 Accuracy 0.3568\n",
      "Epoch 1 Batch 1450 Loss 0.3186 Accuracy 0.3571\n",
      "Epoch 1 Batch 1500 Loss 0.3180 Accuracy 0.3570\n",
      "Epoch 1 Batch 1550 Loss 0.3183 Accuracy 0.3573\n",
      "Epoch 1 Batch 1600 Loss 0.3178 Accuracy 0.3575\n",
      "Epoch 1 Batch 1650 Loss 0.3177 Accuracy 0.3576\n",
      "Epoch 1 Batch 1700 Loss 0.3176 Accuracy 0.3578\n",
      "Epoch 1 Batch 1750 Loss 0.3174 Accuracy 0.3581\n",
      "Epoch 1 Batch 1800 Loss 0.3169 Accuracy 0.3578\n",
      "Epoch 1 Batch 1850 Loss 0.3165 Accuracy 0.3580\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(EPOCHS):\n",
    "  start = time.time()\n",
    "  \n",
    "  train_loss.reset_states()\n",
    "  train_accuracy.reset_states()\n",
    "  \n",
    "  # inp -> peps, tar -> ions\n",
    "  for (batch, (inp, tar)) in enumerate(train_dataset):\n",
    "    train_step(inp, tar)\n",
    "    \n",
    "    if batch % 50 == 0:\n",
    "      print ('Epoch {} Batch {} Loss {:.4f} Accuracy {:.4f}'.format(\n",
    "          epoch + 1, batch, train_loss.result(), train_accuracy.result()))\n",
    "      \n",
    "  if (epoch + 1) % 5 == 0:\n",
    "    ckpt_save_path = ckpt_manager.save()\n",
    "    print ('Saving checkpoint for epoch {} at {}'.format(epoch+1,\n",
    "                                                         ckpt_save_path))\n",
    "    \n",
    "  print ('Epoch {} Loss {:.4f} Accuracy {:.4f}'.format(epoch + 1, \n",
    "                                                train_loss.result(), \n",
    "                                                train_accuracy.result()))\n",
    "\n",
    "  print ('Time taken for 1 epoch: {} secs\\n'.format(time.time() - start))\n",
    "    \n",
    "  test_loss.reset_states()\n",
    "  test_accuracy.reset_states()\n",
    "\n",
    "  # inp -> peps, tar -> ions\n",
    "  for(batch, (inp, tar)) in enumerate(test_dataset):\n",
    "    test_step(inp, tar)\n",
    "\n",
    "    if batch % 100 == 0:\n",
    "      print ('Epoch {} Batch {} Test Loss {:.4f} Test Accuracy {:.4f}'.format(\n",
    "          epoch + 1, batch, test_loss.result(), test_accuracy.result()))\n",
    "\n",
    "  print ('Overall Test Loss {:.4f} Accuracy {:.4f}'.format(test_loss.result(), \n",
    "                                                test_accuracy.result()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluate\n",
    "\n",
    "The following steps are used for evaluation:\n",
    "\n",
    "* Encode the input sentence using the peptide sequence tokenizer (`tokenizer_pep`). Moreover, add the start and end token so the input is equivalent to what the model is trained with. This is the encoder input.\n",
    "* The decoder input is the `start token == tokenizer_ion.vocab_size`.\n",
    "* Calculate the padding masks and the look ahead masks.\n",
    "* The `decoder` then outputs the predictions by looking at the `encoder output` and its own output (self-attention).\n",
    "* Select the last word and calculate the `argmax` of that.\n",
    "* Concatentate the predicted word to the decoder input as pass it to the decoder.\n",
    "* In this approach, the decoder predicts the next word based on the previous words it predicted."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(inp_sentence):\n",
    "  start_token = [tokenizer_pep.vocab_size]\n",
    "  end_token = [tokenizer_pep.vocab_size + 1]\n",
    "  \n",
    "  # inp sentence is portuguese, hence adding the start and end token\n",
    "  inp_sentence = start_token + tokenizer_pep.encode(inp_sentence) + end_token\n",
    "  encoder_input = tf.expand_dims(inp_sentence, 0)\n",
    "  \n",
    "  # as the target is Ions, the first word to the transformer should be the\n",
    "  # ions start token.\n",
    "  decoder_input = [tokenizer_ion.vocab_size]\n",
    "  output = tf.expand_dims(decoder_input, 0)\n",
    "    \n",
    "  for i in range(MAX_LENGTH):\n",
    "    enc_padding_mask, combined_mask, dec_padding_mask = create_masks(\n",
    "        encoder_input, output)\n",
    "  \n",
    "    # predictions.shape == (batch_size, seq_len, vocab_size)\n",
    "    predictions, attention_weights = transformer(encoder_input, \n",
    "                                                 output,\n",
    "                                                 False,\n",
    "                                                 enc_padding_mask,\n",
    "                                                 combined_mask,\n",
    "                                                 dec_padding_mask)\n",
    "    \n",
    "    # select the last word from the seq_len dimension\n",
    "    predictions = predictions[: ,-1:, :]  # (batch_size, 1, vocab_size)\n",
    "\n",
    "    predicted_id = tf.cast(tf.argmax(predictions, axis=-1), tf.int32)\n",
    "    \n",
    "    # return the result if the predicted_id is equal to the end token\n",
    "    if predicted_id == tokenizer_ion.vocab_size+1:\n",
    "      return tf.squeeze(output, axis=0), attention_weights\n",
    "    \n",
    "    # concatentate the predicted_id to the output which is given to the decoder\n",
    "    # as its input.\n",
    "    output = tf.concat([output, predicted_id], axis=-1)\n",
    "\n",
    "  return tf.squeeze(output, axis=0), attention_weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_attention_weights(attention, sentence, result, layer):\n",
    "  fig = plt.figure(figsize=(16, 8))\n",
    "  \n",
    "  sentence = tokenizer_pep.encode(sentence)\n",
    "\n",
    "  attention = tf.squeeze(attention[layer], axis=0)\n",
    "  \n",
    "  for head in range(attention.shape[0]):\n",
    "    ax = fig.add_subplot(2, 4, head+1)\n",
    "    \n",
    "    # plot the attention weights\n",
    "    ax.matshow(attention[head][:-1, :], cmap='viridis')\n",
    "\n",
    "    fontdict = {'fontsize': 10}\n",
    "    \n",
    "    ax.set_xticks(range(len(sentence)+2))\n",
    "    ax.set_yticks(range(len(result)))\n",
    "    \n",
    "    ax.set_ylim(len(result)-1.5, -0.5)\n",
    "        \n",
    "    ax.set_xticklabels(\n",
    "        ['<start>']+[tokenizer_pep.decode([i]) for i in sentence]+['<end>'], \n",
    "        fontdict=fontdict, rotation=90)\n",
    "    \n",
    "    ax.set_yticklabels([tokenizer_ion.decode([i]) for i in result \n",
    "                        if i < tokenizer_ion.vocab_size], \n",
    "                       fontdict=fontdict)\n",
    "    \n",
    "    ax.set_xlabel('Head {}'.format(head+1))\n",
    "  \n",
    "  plt.tight_layout()\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def translate(sentence, plot=''):\n",
    "  result, attention_weights = evaluate(sentence)\n",
    "  \n",
    "  predicted_sentence = tokenizer_ion.decode([i for i in result \n",
    "                                            if i < tokenizer_ion.vocab_size])  \n",
    "\n",
    "  print('Input: {}'.format(sentence))\n",
    "  print('Predicted translation: {}'.format(predicted_sentence))\n",
    "  \n",
    "  if plot:\n",
    "    plot_attention_weights(attention_weights, sentence, result, plot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Translate Peptide Sequences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "translate(\"2 A A A D D G E E P K 22.4 -\", plot='')\n",
    "print('Real Translation: EA KD KF -')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "translate('2 T A G W N I P M G M L Y N R 47.4 -', plot='')\n",
    "print ('Real Translation: ')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "translate('3 P E P T I D E K 52.8 -', plot='')\n",
    "print ('Real Translation: ')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Summary\n",
    "\n",
    "In this notebook, we convert the peptide sequence to ion sequence generation problem to portuguese to english translation problem. The obtained accuracy is about 30%."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Future Work\n",
    "\n",
    "* Insert the charge and energy information in the last dense layer instead of inside sequences\n",
    "* Prediction of internal ions + immonium ions (and their isotopes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
